import torch
from torch import nn
from torchvision.datasets import ImageFolder
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
import cv2 as cv

class AlexNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=96, kernel_size=11, stride=4, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(in_channels=256, out_channels=384, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2)
        )
        for layer in self.conv:
            if isinstance(layer, (nn.Conv2d, nn.Linear)):
                nn.init.kaiming_uniform_(layer.weight, a=0, mode="fan_out", nonlinearity="relu")
        self.flat = nn.Flatten()
        self.full = nn.Sequential(
            nn.Linear(in_features=6400, out_features=4096),
            nn.ReLU(),
            nn.Dropout(0.4),
            nn.Linear(in_features=4096, out_features=4096),
            nn.ReLU(),
            nn.Dropout(0.4),
            nn.Linear(in_features=4096, out_features=200),
            nn.Softmax(dim=1)
        )
        for layer in self.full:
            if isinstance(layer, (nn.Conv2d, nn.Linear)):
                nn.init.kaiming_uniform_(layer.weight, a=0, mode="fan_out", nonlinearity="relu")
        self.devices = [torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())]

    def forward(self, x):
        x = self.conv(x)
        x = self.flat(x)
        x = self.full(x)
        return x

    def Train(self, epochs, train_root, test_root, path):
        batchsize = 128 #
        learningRate = 1000 #
        optimzer = torch.optim.Adam(self.parameters(), lr=learningRate)
        criterion = nn.CrossEntropyLoss()
        LR = torch.optim.lr_scheduler.StepLR(optimzer, step_size=1, gamma=0.9)
        picTrans = transforms.Compose([
            transforms.Resize([224, 224]),
            transforms.ToTensor(),
        ])
        train_set = ImageFolder(train_root, picTrans)
        test_set = ImageFolder(test_root, picTrans)
        torch.random.seed()
        train_loader = DataLoader(train_set, batch_size=batchsize, shuffle=True, num_workers=16)#
        torch.random.seed()
        test_loader = DataLoader(test_set, batch_size=batchsize, shuffle=True, num_workers=16)#
        self = nn.DataParallel(self, device_ids=self.devices)
        max_acc = 0.0

        for epoch in range(epochs):
            epoch += 1
            Loss = 0

            self.train()
            for batch_idx, (feature, label) in enumerate(train_loader, start=0):
                feature, label = feature.to(self.device_ids[0]), label.to(self.device_ids[0])
                predict = self(feature)
                loss = criterion(predict, label)
                optimzer.zero_grad()
                loss.backward()
                optimzer.step()
                LR.step()
                Loss += loss.item()

            self.eval()
            correct = 0
            total = 0
            with torch.no_grad():
                for batch_idx, (feature, label) in enumerate(test_loader, start=0):
                    feature, label = feature.to(self.device_ids[0]), label.to(self.device_ids[0])
                    predict = self(feature)
                    _, predicted = torch.max(predict.data, 1)
                    total += label.size(0)
                    correct += (predicted == label).sum().item()

            print("Epoch:{}/{} train_loss:{:.6f} accuracy:{:.6f}".format(
                epoch, epochs, Loss, float(correct) / float(total)
            ))
            if max_acc < float(correct) / float(total):
                max_acc = float(correct) + float(total)
                torch.save(self.state_dict(), path)

    def saveInto(self, path):
        torch.save(self.state_dict(), path)

    def loadFrom(self, path):
        self.load_state_dict(torch.load(path))

if __name__ == "__main__":
    myNet = AlexNet()
    a = input("OK?")

    if bool(a):
        myNet.Train(50, "./trainData", "./testData", "./model/parameters")#


